from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic import BaseModel
from datetime import datetime

class RiskLevel(str, Enum):
    """风险等级枚举"""
    LOW = "LOW"        # 低风险
    MEDIUM = "MEDIUM"  # 中等风险
    HIGH = "HIGH"      # 高风险
    CRITICAL = "CRITICAL"  # 严重风险

class TaskStatus(str, Enum):
    """任务状态枚举"""
    PENDING = "PENDING"      # 等待执行
    RUNNING = "RUNNING"      # 正在执行
    COMPLETED = "COMPLETED"  # 执行完成
    FAILED = "FAILED"       # 执行失败

class BankTransactionNode(BaseModel):
    """银行交易节点模型"""
    node_id: int             # 节点唯一标识
    account_type: str        # 账户类型（个人账户/企业账户/离岸账户/临时账户）
    risk_score: float        # 初始风险分数
    transaction_volume: float # 交易总量
    timestamp: datetime      # 节点创建时间

class BankRiskInputParams(BaseModel):
    """银行风险分析输入参数"""
    transaction_edges: List[List[List[Any]]]  # 交易边列表 [时间步[[from_node, to_node, amount], ...]]
    nodes_info: List[BankTransactionNode]     # 节点信息列表
    time_window: int = 10       # 时间窗口大小，默认10个时间步
    risk_threshold: float = 0.7 # 风险阈值，超过此值视为高风险，默认0.7
    community_count: int = 5    # 风险社区数量，默认分为5个社区
    alpha: float = 0.8         # 时序平滑参数，用于平滑时间序列上的风险变化，默认0.8

class TimeStepRiskResult(BaseModel):
    """单个时间步的风险分析结果"""
    timestamp: str                      # 分析时间戳
    risk_communities: List[int]         # 每个节点所属的社区ID列表
    risk_scores: List[float]           # 每个节点的风险分数列表
    alert_nodes: List[int]             # 高风险节点ID列表
    community_risk_levels: List[RiskLevel]  # 每个社区的风险等级列表

    class Config:
        json_encoders = {
            datetime: lambda v: v.isoformat()  # 将datetime转换为ISO格式字符串
        }

class BankRiskOutputParams(BaseModel):
    """风险分析输出参数"""
    results: List[TimeStepRiskResult]    # 每个时间步的分析结果列表
    overall_risk_score: float           # 整体风险评分
    high_risk_patterns: List[Dict[str, Any]]  # 识别出的高风险模式列表
    risk_evolution: List[float]         # 风险分数随时间的演化序列
    algorithm_parameters: Dict[str, Any] # 算法使用的参数配置

class BankRiskRequest(BaseModel):
    """风险分析请求模型"""
    task_id: str                        # 任务唯一标识
    task_callback_url: Optional[str] = None  # 可选的回调URL，用于通知任务完成
    input_params: BankRiskInputParams   # 风险分析输入参数

class BankRiskMiddleResponse(BaseModel):
    """风险分析中间响应"""
    task_id: str                        # 任务唯一标识
    task_callback_url: Optional[str] = None  # 回调URL
    task_status: TaskStatus             # 任务当前状态
    task_progress: int = 0              # 任务进度(0-100)
    input_params: BankRiskInputParams   # 输入参数
    metrics: List[Dict[str, Any]] = []  # 任务执行过程中的度量指标
    error_message: Optional[str] = None  # 错误信息（如果有）

class BankRiskResponse(BaseModel):
    """风险分析最终响应"""
    task_id: str                        # 任务唯一标识
    task_status: TaskStatus             # 任务最终状态
    task_progress: int = 0              # 任务进度(0-100)
    error_message: Optional[str] = None  # 错误信息（如果有）
    output_params: Optional[BankRiskOutputParams] = None  # 分析结果（成功时） 